#
from typing import Dict
import numpy as np

class MimoTx(object):
    def __init__(self):
        self.name = 'study.aic.mimo_tx.MimoTx'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'MIMO发射端 v0.0.2')
        MimoTx.ofdm_tx()

    @staticmethod
    def ofdm_tx(params:Dict = {}) -> None:
        # 参数设置
        num_subcarriers = 64  # 子载波数量
        num_symbols = 100  # 符号数量
        cp_length = 16  # 循环前缀长度
        carrier_frequency = 1e6  # 载波频率
        sampling_rate = 10e6  # 采样率
        # 生成随机数据符号
        data_symbols = np.random.randint(0, 2, size=(num_symbols, num_subcarriers)) * 2 - 1
        # IFFT 变换
        ifft_data = np.fft.ifft(data_symbols, axis=1)
        # 添加循环前缀
        ifft_data_with_cp = np.concatenate((ifft_data[:, -cp_length:], ifft_data), axis=1)
        # 并串转换
        serial_data = ifft_data_with_cp.flatten()
        print(f'### 0 len={len(serial_data)};')
        print(f'### 1 serial_data: {serial_data.shape}; sampling_rate: {sampling_rate}; up: {len(serial_data) / sampling_rate}; step: {1 / sampling_rate};')
        # 生成载波信号
        # t = np.arange(0, len(serial_data) / sampling_rate, 1 / sampling_rate) # 模型生成的代码有问题，会多出一个数据点
        t = np.arange(len(serial_data))*(1/sampling_rate)
        print(f'### 2 t: {t.shape}; \n{t};')
        carrier_signal = np.cos(2 * np.pi * carrier_frequency * t)
        # 调制到载波上
        print(f'serial_data: {serial_data.shape}; carrier_signal: {carrier_signal.shape};')
        modulated_signal = serial_data * carrier_signal
        # 输出结果
        print("Generated OFDM signal with {} subcarriers and {} symbols, modulated onto a carrier at {} Hz and sampled at {} Hz.".format(num_subcarriers, num_symbols, carrier_frequency, sampling_rate))
        params['modulated_signal'] = modulated_signal
        MimoTx.ofdm_rx(params=params)

    @staticmethod
    def ofdm_rx(params:Dict = {}) -> None:
        # 接收端代码
        # 假设已经接收到了 IQ 采样数据
        # 这里我们使用随机生成的数据作为示例
        sampling_rate = 10e6  # 采样率
        carrier_frequency = 1e6  # 载波频率
        num_symbols = 100  # 符号数
        num_subcarriers = 64  # 子载波数
        cp_length = 16  # 循环前缀长度

        # 生成随机的 IQ 采样数据
        # received_iq_samples = np.random.randn(num_symbols * (num_subcarriers + cp_length)) + 1j * np.random.randn(num_symbols * (num_subcarriers + cp_length))
        modulated_signal = params['modulated_signal']
        received_iq_samples = modulated_signal[::2] + 1j * modulated_signal[1::2]

        # 解调
        # demodulated_signal = received_iq_samples.real * np.cos(2 * np.pi * carrier_frequency * np.arange(0, len(received_iq_samples.real) / sampling_rate, 1 / sampling_rate)) - received_iq_samples.imag * np.sin(2 * np.pi * carrier_frequency * np.arange(0, len(received_iq_samples.imag) / sampling_rate, 1 / sampling_rate))
        t = np.arange(len(received_iq_samples.imag) / sampling_rate) * (1 / sampling_rate)
        demodulated_signal = received_iq_samples.real * np.cos(2 * np.pi * carrier_frequency * t) - \
                    received_iq_samples.imag * np.sin(2 * np.pi * carrier_frequency * t)

        # 串并转换
        received_ifft_data_with_cp = demodulated_signal.reshape(num_symbols, -1)

        # 去除循环前缀
        received_ifft_data = received_ifft_data_with_cp[:, cp_length:]

        # FFT 变换
        received_data_symbols = np.fft.fft(received_ifft_data, axis=1)

        # 数据同步：假设使用自相关方法找到符号起始位置
        # 这里需要确保 received_serial_data 已经被正确定义或赋值
        received_serial_data = demodulated_signal  # 假设 demodulated_signal 就是 received_serial_data

        sync_position = np.argmax(np.correlate(received_serial_data, received_serial_data, mode='same')) - len(received_serial_data) // 2
        received_ifft_data = received_ifft_data[sync_position:]

        # 信道估计：假设使用导频信号进行估计
        pilot_symbols = received_data_symbols[::cp_length, :cp_length]  # 假设每10个符号有10个导频符号
        channel_estimate = np.mean(pilot_symbols, axis=0)
        equalized_data_symbols = received_data_symbols / channel_estimate

        # 均衡：假设使用迫零均衡算法
        # # 将 channel_estimate 扩展为与 received_data_symbols 形状相匹配的数组
        # channel_estimate_expanded = np.tile(channel_estimate, (num_symbols, 1))
        # equalized_data_symbols = received_data_symbols / channel_estimate_expanded

        # # 将 channel_estimate 扩展为与 received_data_symbols 形状相匹配的数组
        # channel_estimate_expanded = np.tile(channel_estimate, (1, num_subcarriers))
        # print(f'### received_data_symbols: {received_data_symbols.shape}; channel_estimate_expanded: {channel_estimate_expanded.shape}; channel_estimate: {channel_estimate.shape};')
        # equalized_data_symbols = received_data_symbols / channel_estimate_expanded

        # 输出结果
        print("Received OFDM signal with {} subcarriers and {} symbols after synchronization, channel estimation, and equalization.".format(num_subcarriers, num_symbols))





































    @staticmethod
    def ofdm_raw(params:Dict = {}) -> None:
        # 参数设置
        num_subcarriers = 64  # 子载波数量
        num_symbols = 100  # 符号数量
        cp_length = 16  # 循环前缀长度
        # 生成随机数据符号
        data_symbols = np.random.randint(0, 2, size=(num_symbols, num_subcarriers)) * 2 - 1
        print(f'data_symbols: {data_symbols.shape};\n{data_symbols};')
        # IFFT 变换
        ifft_data = np.fft.ifft(data_symbols, axis=1)
        print(f'ifft_data: {ifft_data.shape};')
        # 添加循环前缀
        ifft_data_with_cp = np.concatenate((ifft_data[:, -cp_length:], ifft_data), axis=1)
        print(f'ifft_data_with_cp: {ifft_data_with_cp.shape};')
        # 并串转换
        serial_data = ifft_data_with_cp.flatten()
        print(f'serial_data: {serial_data.shape};')
        # 输出结果
        print("Generated OFDM signal with {} subcarriers and {} symbols.".format(num_subcarriers, num_symbols))        